#ifndef __JNS_io_hpp__
#define __JNS_io_hpp__

#ifndef _FILE_OFFSET_BITS
#define _FILE_OFFSET_BITS 64
#endif

#ifndef _LARGEFILE64_SOURCE
#define _LARGEFILE64_SOURCE
#endif

/*
  (c) James N. Simone and Fermilab
*/

#include <sys/types.h>
#include <unistd.h>
#include <vector>
#include <io_uni.h>

#include <Wtimer.hpp>

// === BEGIN external QCD types ===
//
// Note: external QCD types may be mapped in more than one
// mdp type. Type safety in mdp is incomplete.

template<typename R>
struct io_complex
{
  typedef R base_type;
  typedef generic_field<mdp_complex> field_type;

  // extract value to f(s)
  // Assumption: generic_field<mdp_complex> f has a single element
  inline void get ( field_type& f, site const& s ) const
  {
    f ( s ).re = real;
    f ( s ).im = imag;
  }

  // set value from f(s)
  // mdp bug: inline void set ( field_type const& f, site const& s )
  inline void set ( field_type& f, site const& s )
  {
    real = f ( s ).re;
    imag = f ( s ).im;
  }

  R real;
  R imag;
};

template<typename R>
struct io_gl3_matrix
{
  typedef R base_type;
  typedef mdp_matrix_field field_type;

  static int const nC = 3;

  // extract value to f(s)
  // Assumption: mdp_matrix_field f has 3x3 complex elements
  inline void get ( field_type& f, site const& s ) const
  {
    for ( int a = 0; a < nC; ++a )
      for ( int b = 0; b < nC; ++b )
	{
	  f ( s, a, b ).re = e [a][b].real;
	  f ( s, a, b ).im = e [a][b].imag;
	}
  }

  // set value from f(s)
  // mdp bug: inline void set ( field_type const& f, site const& s )
  inline void set ( field_type& f, site const& s )
  {
    for ( int a = 0; a < nC; ++a )
      for ( int b = 0; b < nC; ++b )
	{
	  e [a][b].real = f ( s, a, b ).re;
	  e [a][b].imag = f ( s, a, b ).im;
	}
  }

  io_complex<R> e [nC][nC];
};

template<typename R>
struct gl3_links
{
  typedef R base_type;
  typedef gauge_field field_type;

  static int const nDir = 4;
  static int const nC   = 3;

  // extract value to f(s)
  // Assumption: gauge_field f has 4x3x3 complex elements
  inline void get ( field_type& f, site const& s ) const
  {
    // direction: T X Y Z
    //  internal: 0 1 2 3
    //  file    : 3 0 1 2
    static int const mapdir [nDir] = { 3, 0, 1, 2 };

    // store site
    for ( int mu = 0; mu < nDir; ++mu )
      {
	int nu = mapdir [mu];

	for ( int a = 0; a < nC; ++a )
	  for ( int b = 0; b < nC; ++b )
	    {
	      f ( s, mu, a, b ).re = link [nu][a][b].real;
	      f ( s, mu, a, b ).im = link [nu][a][b].imag;
	    }
      }
  }

  // set value from f(s)
  // mdp bug: inline void set ( field_type const& f, site const& s )
  inline void set ( field_type& f, site const& s )
  {
    // direction: T X Y Z
    //  internal: 0 1 2 3
    //  file    : 3 0 1 2
    static int const mapdir [nDir] = { 3, 0, 1, 2 };

    for ( int mu = 0; mu < nDir; ++mu )
      {
	int nu = mapdir [mu];

	for ( int a = 0; a < nC; ++a )
	  for ( int b = 0; b < nC; ++b )
	    {
	      link [nu][a][b].real = f ( s, mu, a, b ).re;
	      link [nu][a][b].imag = f ( s, mu, a, b ).im;
	    }
      }
  }

  io_complex<R> link [nDir][nC][nC];
};

template<typename R>
struct Wilson_sc_quark
{
  typedef R base_type;
  typedef fermi_field field_type;

  static int const nS = 4;
  static int const nC = 3;

  // extract value to f(s)
  // Assumption: fermi_field f has 4x3 complex elements
  inline void get ( field_type& f, site const& s ) const
  {
    for ( int a = 0; a < nS; ++a )
      for ( int b = 0; b < nC; ++b )
	{
	  f ( s, a, b ).re = q [a][b].real;
	  f ( s, a, b ).im = q [a][b].imag;
	}
  }

  // set value from f(s)
  // mdp bug: inline void set ( field_type const& f, site const& s )
  inline void set ( field_type& f, site const& s )
  {
    for ( int a = 0; a < nS; ++a )
      for ( int b = 0; b < nC; ++b )
	{
	  q [a][b].real = f ( s, a, b ).re;
	  q [a][b].imag = f ( s, a, b ).im;
	}
  }

  io_complex<R> q [nS][nC];
};

// === END external QCD types ===


template<typename T>
struct ioBuffer
{
  T*   base;
  long nElem;

  ioBuffer ( unsigned len ) : nElem ( len )
  {
    base = new T [nElem];
  }

  ~ioBuffer ( )
  {
    delete base;
  }
};

int openFile ( IO_filehand& fh, IO_uni_header& header,
	       char const* filename,
	       char const* mode, int ioNode = 0 )
{
  int selfNode = mpi.me ( );
  int retVal = 0;

  if ( selfNode == ioNode )
    {
      // open binary and metadata files
      retVal = IO_open ( &fh, filename, mode );
      if ( retVal )
	return retVal;

      string smode ( mode );
      if ( smode != "w" )
	{
	  if ( smode == "a+" )
	    {
	      // seek to beginning of file before header read
	      int status;
	      if ( ( status = fseeko64 ( fh.data_fp, 0L, SEEK_SET ) ) )
		{
		  cerr << "Cannot seek file (status = "
			    << status << ")" << endl;
		  exit ( 1 );
		}
	    }

	  // read binary header
	  retVal = IO_read_header ( &fh, &header );
	}
    }

  return retVal;
}

int closeFile ( IO_filehand& fh, int ioNode = 0 )
{
  int selfNode = mpi.me ( );
  int retVal = 0;

  if ( selfNode == ioNode )
    {
      retVal = IO_close ( &fh );
    }
  return retVal;
}

// return updated siteCnt, offset, x0, y0, z0, t0
long lookAhead ( long sitesMax, long& siteCnt, mdp_lattice& lattice,
		 vector<site>& offset,
		 long& t0, long& z0, long& y0, long& x0 )
{
  offset.clear ( ); // empty
  siteCnt = 0;
  site s ( lattice );
  long t = 0, z = 0, y = 0, x = 0;
  bool first = true;

  long destNode = on_which_process  ( lattice, t0, x0, y0, z0 );

  bool cont = true;
  for ( t = 0; cont && t < lattice.size (TIME); ++t )
    for ( z = 0; cont && z < lattice.size (SPACE_Z); ++z )
      for ( y = 0; cont && y < lattice.size (SPACE_Y); ++y )
	for ( x = 0; cont && x < lattice.size (SPACE_X); ++x )
	  {
	    if ( first )
	      {
		// jump back into loop
		first = false;
		t = t0;
		z = z0;
		y = y0;
		x = x0;
	      }

	    if ( destNode == on_which_process  ( lattice, t, x, y, z )
		 && siteCnt < sitesMax )
	      {
		if ( destNode == mpi.me ( ) )
		  {
		    // add site
		    s.set ( t, x, y, z );
		    offset.push_back ( s );
		  }

		// bump count
		++siteCnt;
	      }
	    else
	      {
		// save last coord
		t0 = t;
		z0 = z;
		y0 = y;
		x0 = x;

		// exit look ahead
		cont = false;
	      }
	  }

  if ( cont )
    {
      // save last coord
      t0 = t;
      z0 = z;
      y0 = y;
      x0 = x;
    }

  return destNode;
}

template<typename T>
void receiveSites ( ioBuffer<T>& buf, long nSites,
		    int destNode, int ioNode = 0 )
{
  int selfNode = mpi.me ( );

  if ( selfNode == destNode && selfNode != ioNode )
    {
      mpi.get ( buf.base, nSites, ioNode ); // blocking
    }
}

template<typename T>
void sendSites ( ioBuffer<T>& buf, long nSites,
		 int srcNode, int ioNode = 0 )
{
  int selfNode = mpi.me ( );

  if ( selfNode == srcNode && selfNode != ioNode )
    {
      mpi.put ( buf.base, nSites, ioNode ); // blocking
    }
}

template<typename T>
int readSites ( IO_filehand& fh, IO_uni_header& header,
		ioBuffer<T>& buf, long nSites,
		int destNode, int ioNode = 0 )
{
  int retVal = 0;
  int selfNode = mpi.me ( );

  if ( selfNode == ioNode )
    {
      retVal = IO_read_site_data ( &fh, &header, nSites, buf.base );

      if ( selfNode != destNode )
	{
	  mpi.put ( buf.base, nSites, destNode ); // blocking
	}
    }

  return retVal;
}

template<typename T>
int writeSites ( IO_filehand& fh, IO_uni_header& header,
		 ioBuffer<T>& buf, long nSites,
		 int srcNode, int ioNode = 0 )
{
  int retVal = 0;
  int selfNode = mpi.me ( );

  if ( selfNode == ioNode )
    {
      if ( selfNode != srcNode )
	{
	  mpi.get ( buf.base, nSites, srcNode ); // blocking
	}

      retVal = IO_write_site_data ( &fh, &header, nSites, buf.base );
    }

  return retVal;
}

template<typename T>
void scatterSites ( typename T::field_type& f,
		    vector<site> const& offset,
		    int destNode, ioBuffer<T> const& buf )
{
  int selfNode = mpi.me ( );

  if ( selfNode == destNode )
    {
      for ( unsigned j = 0; j < offset.size ( ); ++j )
	{
	  // store site in f
	  buf.base [j].get ( f, offset [j] );
	}
    }
}

template<typename T>
void gatherSites ( typename T::field_type& f,
		   vector<site> const& offset,
		   int srcNode, ioBuffer<T>&  buf )
{
  int selfNode = mpi.me ( );

  if ( selfNode == srcNode )
    {
      for ( unsigned j = 0; j < offset.size ( ); ++j )
	{
	  // get site
	  buf.base [j].set ( f, offset [j] );
	}
    }
}

// example: JNS_read<gl3_links<float> >
// record number is zero based
template<typename T>
void JNS_read ( typename T::field_type& f, char const* filename,
		unsigned record,
		int bufSize, int ioNode = 0 )
{
  Wtimer timer;
  timer.start ( );

  int selfNode = mpi.me ( );

  int const sitesPerRead = bufSize;
  ioBuffer<T> buf ( sitesPerRead );
  vector<site> offset;
  offset.reserve ( sitesPerRead );

  int status;
  IO_filehand fh;
  IO_uni_header header;

  status = openFile ( fh, header, filename, "r" );
  if ( status )
    {
      cerr << "ERROR open failed: " << filename << endl;
      exit ( 1 );
    }

  // sanity checks
  if ( selfNode == ioNode )
    {
      bool quit = false;

     if ( header.size_of_element != sizeof(T::base_type) )
        {
	  quit = true;
	  cerr << "wrong size-of-element" << endl;
	}
     if ( header.elements_per_site != sizeof(T)/sizeof(T::base_type) )
       {
	 quit = true;
	 cerr << "wrong elements-per-site" << endl;
       }
      if ( header.dim [0] != f.lattice ( ).size (SPACE_X) )
	{
	  quit = true;
	  cerr << "lattice mismatch X: "
		    << header.dim [0]
		    << " != "
		    << f.lattice ( ).size (SPACE_X)
		    << endl;
	}
      if ( header.dim [1] != f.lattice ( ).size (SPACE_Y) )
	{
	  quit = true;
	  cerr << "lattice mismatch Y: "
		    << header.dim [1]
		    << " != "
		    << f.lattice ( ).size (SPACE_Y)
		    << endl;
	}
      if ( header.dim [2] != f.lattice ( ).size (SPACE_Z) )
	{
	  quit = true;
	  cerr << "lattice mismatch Z: "
		    << header.dim [2]
		    << " != "
		    << f.lattice ( ).size (SPACE_Z)
		    << endl;
	}
      if ( header.dim [3] != f.lattice ( ).size (TIME) )
	{
	  quit = true;
	  cerr << "lattice mismatch T: "
		    << header.dim [3]
		    << " != "
		    << f.lattice ( ).size (TIME)
		    << endl;
	}
      if ( quit )
	exit ( 1 );
    }

  long totalSites =
    f.lattice ( ).size (TIME) *
    f.lattice ( ).size (SPACE_Z) *
    f.lattice ( ).size (SPACE_Y) *
    f.lattice ( ).size (SPACE_X);

  if ( selfNode == ioNode )
    {
      off_t foffset = totalSites;
      foffset *= header.elements_per_site;
      foffset *= header.size_of_element;
      foffset *= record;

      // seek relative to current position
      if ( ( status = fseeko64 ( fh.data_fp, foffset, SEEK_CUR ) ) )
	{
	  cerr << "Cannot seek file (status = "
		    << status << ")" << endl;
	  exit ( 1 );
	}
    }

  long Sites, siteCnt;
  int nReads = 0;
  long t = 0, z = 0, y = 0, x = 0;
  for ( Sites = 0; Sites < totalSites; Sites += siteCnt )
    {
      long destNode =lookAhead ( sitesPerRead,
				 siteCnt, f.lattice ( ),
				 offset,
				 t, z, y, x );

      receiveSites ( buf, siteCnt, destNode );

      readSites ( fh, header, buf, siteCnt, destNode );

      scatterSites ( f, offset, destNode, buf );

      ++nReads;
    }

  mpi.barrier ( );

  f.update ( );

  closeFile ( fh );

  double elapsed = timer.stop ( );
  if ( selfNode == ioNode )
    {
      cout << "JNS_read: time " << elapsed << " secs" << endl;
      cout << "disk reads: " << nReads << " (" << Sites << " sites )" << endl;
    }
}

// TODO:
template<typename T>
void JNS_write ( typename T::field_type& f, char const* filename,
		 unsigned record,
		 int bufSize, int ioNode = 0 )
{
  Wtimer timer;
  timer.start ( );

  int selfNode = mpi.me ( );

  int const sitesPerWrite = bufSize;
  ioBuffer<T> buf ( sitesPerWrite );
  vector<site> offset;
  offset.reserve ( sitesPerWrite );

  int status;
  IO_filehand fh;
  IO_uni_header header;

  char const* mode = ( record == 0 ) ? "w" : "a+";
  // mode "w"  truncate/create open write only
  // mode "a+" open write/read position at EOF
  status = openFile ( fh, header, filename, mode );
  if ( status )
    {
      cerr << "ERROR open failed: " << filename << endl;
      exit ( 1 );
    }

  if ( selfNode == ioNode )
    {
      if ( record == 0 )
	{
	  // prepare header
	  static int const nDir = 4;
	  int dim [nDir] = {
	    f.lattice ( ).size (SPACE_X),
	    f.lattice ( ).size (SPACE_Y),
	    f.lattice ( ).size (SPACE_Z),
	    f.lattice ( ).size (TIME),
	  };

	  IO_set_header_for_write ( &fh, &header, sizeof(T::base_type),
				    sizeof(T) / sizeof(T::base_type),
				    &dim [0], IO_NATURAL );

	  // write header
	  IO_write_header ( &fh, &header );
	}

      // sanity checks
      bool quit = false;

     if ( header.size_of_element != sizeof(T::base_type) )
        {
	  quit = true;
	  cerr << "wrong size-of-element" << endl;
	}
     if ( header.elements_per_site != sizeof(T)/sizeof(T::base_type) )
       {
	 quit = true;
	 cerr << "wrong elements-per-site" << endl;
       }
      if ( header.dim [0] != f.lattice ( ).size (SPACE_X) )
	{
	  quit = true;
	  cerr << "lattice mismatch X: "
		    << header.dim [0]
		    << " != "
		    << f.lattice ( ).size (SPACE_X)
		    << endl;
	}
      if ( header.dim [1] != f.lattice ( ).size (SPACE_Y) )
	{
	  quit = true;
	  cerr << "lattice mismatch Y: "
		    << header.dim [1]
		    << " != "
		    << f.lattice ( ).size (SPACE_Y)
		    << endl;
	}
      if ( header.dim [2] != f.lattice ( ).size (SPACE_Z) )
	{
	  quit = true;
	  cerr << "lattice mismatch Z: "
		    << header.dim [2]
		    << " != "
		    << f.lattice ( ).size (SPACE_Z)
		    << endl;
	}
      if ( header.dim [3] != f.lattice ( ).size (TIME) )
	{
	  quit = true;
	  cerr << "lattice mismatch T: "
		    << header.dim [3]
		    << " != "
		    << f.lattice ( ).size (TIME)
		    << endl;
	}
      if ( quit )
	exit ( 1 );
    }

  long totalSites =
    f.lattice ( ).size (TIME) *
    f.lattice ( ).size (SPACE_Z) *
    f.lattice ( ).size (SPACE_Y) *
    f.lattice ( ).size (SPACE_X);

  if ( selfNode == ioNode )
    {
      off_t foffset = totalSites;
      foffset *= header.elements_per_site;
      foffset *= header.size_of_element;
      foffset *= record;

      // seek relative to current position
      if ( ( status = fseeko64 ( fh.data_fp, foffset, SEEK_CUR ) ) )
	{
	  cerr << "Cannot seek file (status = "
		    << status << ")" << endl;
	  exit ( 1 );
	}
    }

  long Sites, siteCnt;
  int nWrites = 0;
  long t = 0, z = 0, y = 0, x = 0;
  long lastSrcNode = -1;
  for ( Sites = 0; Sites < totalSites; Sites += siteCnt )
    {
      long srcNode =lookAhead ( sitesPerWrite,
				siteCnt, f.lattice ( ),
				offset,
				t, z, y, x );

      gatherSites ( f, offset, srcNode, buf );

      sendSites ( buf, siteCnt, srcNode );

      writeSites ( fh, header, buf, siteCnt, srcNode );

      ++nWrites;

      if ( srcNode != lastSrcNode )
	{
	  // barrier when sending node changes
	  mpi.barrier ( );
	  lastSrcNode = srcNode;
	}
    }
  mpi.barrier ( );

  closeFile ( fh );

  double elapsed = timer.stop ( );
  if ( selfNode == ioNode )
    {
      cout << "JNS_write: time " << elapsed << " secs" << endl;
      cout << "disk writes: " << nWrites << " (" << Sites << " sites )"
		<< endl;
    }
}

#endif











